package com.ztesoft.zcm.cmdb.util.remote;

import com.jcraft.jsch.ChannelExec;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.ztesoft.zcm.cmdb.util.CmdbExceptionErrorCode;
import com.ztesoft.zsmart.core.exception.BaseAppException;
import com.ztesoft.zsmart.core.log.ZSmartLogger;
import com.ztesoft.zsmart.zcm.core.exception.ExceptionPublisher;

import java.io.ByteArrayOutputStream;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;


/**
 * 提供exec功能的远程客户端
 *
 * @author li.peilong
 * @date 2019/05/24
 **/
public class ExecRemoteClient extends RemoteClient {
    private static final ZSmartLogger LOGGER = ZSmartLogger.getLogger(ExecRemoteClient.class);
    private Session session;
    public ExecRemoteClient(HostAccessInfo accessInfo) {
        super(accessInfo);
    }

    /**
     * 建立连接
     */
    public void connect() throws BaseAppException {
        try {
            this.session = openSession();
        }
        catch (Exception e) {
            this.close();
            ExceptionPublisher.publish(e, CmdbExceptionErrorCode.REMOTE_CONNECT_TO_HOST_FAILED, this.getHostAccessInfo().getHost());
        }
    }
    @Override
    public void close() {
        try {
            if (this.session != null) {
                this.session.disconnect();
                this.session = null;
            }
        }
        catch (Exception e) {
            LOGGER.warn("Exception when closing session", e);
        }
    }

    @Override
    public RemoteResult exec(String cmd, String... patterns) throws BaseAppException {
        return this.exec(cmd, 3, patterns);
    }

    @Override
    public RemoteResult exec(String cmd, int timeout, String... patterns) throws BaseAppException {
        final ByteArrayOutputStream out = new ByteArrayOutputStream();
        final ByteArrayOutputStream errout = new ByteArrayOutputStream();
        final ChannelExec channel;
        try {
            channel = (ChannelExec) session.openChannel("exec");
            channel.setCommand(cmd);
            channel.setOutputStream(out);
            channel.setExtOutputStream(out);
            channel.setErrStream(errout);
            channel.connect();
            // wait for it to finish
            CountDownLatch latch = new CountDownLatch(1);
            pool.execute(() -> {
                while (true) {
                    if (channel.isClosed()) {
                        latch.countDown();
                        break;
                    }
                }
            });
            RemoteResult result = new RemoteResult();

            boolean success = latch.await(timeout, TimeUnit.SECONDS);
            if (success) {
                result.setExitCode(channel.getExitStatus());
                result.setOutput(out.toString());
                result.setError(errout.toString());
                return result;
            }
            else {
                if (channel != null) {
                    channel.disconnect();
                }
                ExceptionPublisher.publish(CmdbExceptionErrorCode.REMOTE_EXEC_COMMAND_TIMEOUT, cmd);
            }
        }
        catch (InterruptedException | JSchException e) {
            ExceptionPublisher.publish(e, CmdbExceptionErrorCode.REMOTE_EXEC_COMMAND_FAILED, cmd);
        }
        return null;
    }
}
